module Misc where

import Types
import Printer(printExpr, printStmt)


isNumeric :: Type -> Bool
isNumeric t = t `elem` [Int, Int2, Int3, Int4, Float, Float2, Float3, Float4]


-- join types into a common super-type
join :: Type -> Type -> Type
join t1 t2 | t1 == t2 = t1
join NoType t = t
join t NoType = t
join Int t | t `elem` [Int2, Int3, Int4, Float, Float2, Float3, Float4] = t
join t Int | t `elem` [Int2, Int3, Int4, Float, Float2, Float3, Float4] = t
join Float t | t `elem` [Float2, Float3, Float4] = t
join t Float | t `elem` [Float2, Float3, Float4] = t
join Int2 Float2 = Float2
join Float2 Int2 = Float2
join Int3 Float3 = Float3
join Float3 Int3 = Float3
join Int4 Float4 = Float4
join Float4 Int4 = Float4
join t1 t2 | otherwise = error $ "incompatible types: " ++ show t1 ++ " " ++ show t2


getFunType :: Function -> FunType
getFunType (Function _ retType args _) = FunType [t | (_, t) <- args] retType


-- cast expression into a desired type
cast :: Type -> Type -> Expr -> Expr
cast _ NoType e = e
cast t1 t2 e | t1 == t2 = e
cast Int Float e = e -- this cast is implicit in cuda
cast Int t e | t `elem` [Int2, Int3, Int4, Float2, Float3, Float4] = Cast t e
cast Float t e | t `elem` [Float2, Float3, Float4] = Cast t e
cast Int2 Float2 e = Cast Float2 e
cast Int3 Float3 e = Cast Float3 e
cast Int4 Float4 e = Cast Float4 e
cast t1 t2 e = error $ "cannot cast " ++ show t1 ++ " into " ++ show t2
                ++ " in " ++ printExpr e


-- remove redundant nested statement groups
removeGroups :: Stmt -> Stmt
removeGroups (Branch c s1 s2) = Branch c (removeGroups s1) (removeGroups s2)
removeGroups (While c s) = While c (removeGroups s)
removeGroups (For c u s) = For c u (removeGroups s)
removeGroups (FunDef (Function n t a b)) = FunDef $ Function n t a (removeGroups b)
removeGroups (Group stmts) = Group (r stmts) where
    r [] = []
    r (Group s: rest) = r (s ++ rest)
    r (x: rest) = removeGroups x : r rest
removeGroups s = s


-- find the (joint) return type of a group of statements
getReturnType :: Stmt -> Type
getReturnType (Group stmts) = foldr join NoType (map getReturnType stmts)
getReturnType (Return _ t) = t
getReturnType (Branch _ s1 s2) = join (getReturnType s1) (getReturnType s2)
getReturnType (While _ s) = getReturnType s
getReturnType (For _ _ s) = getReturnType s
getReturnType _ = NoType


-- cast return types in a statement group
castReturnTypes :: Type -> Stmt -> Stmt  
castReturnTypes t (Group stmts) = Group (map (castReturnTypes t) stmts)
castReturnTypes t (Return e t1) = Return (cast t1 t e) t
castReturnTypes t (Branch c s1 s2) = Branch c s1' s2' where
    s1' = castReturnTypes t s1
    s2' = castReturnTypes t s2
castReturnTypes t (While c s) = While c (castReturnTypes t s)
castReturnTypes t (For c u s) = For c u (castReturnTypes t s)
castReturnTypes t s = s


-- enumerate variables used in statement (duplicates possible)
varsExpr :: Expr -> [String]
varsExpr (V x) = [x]
varsExpr (Func _ args) = concat $ map varsExpr args
varsExpr (Op _ e1 e2) = varsExpr e1 ++ varsExpr e2
varsExpr (Cast _ e) = varsExpr e
varsExpr (Dot e _) = varsExpr e
varsExpr _ = []

varsStmt :: Stmt -> [String]
varsStmt (Group s) = concat $ map varsStmt s
varsStmt (Assign x e) = x : varsExpr e
varsStmt (Branch c s1 s2) = varsExpr c ++ varsStmt s1 ++ varsStmt s2
varsStmt (While c s) = varsExpr c ++ varsStmt s
varsStmt (For c u s) = varsExpr c ++ varsStmt u ++ varsStmt s
varsStmt (Return e _) = varsExpr e
varsStmt _ = []

